import os
import glob
import argparse
from gmssl import *

def test_sm4_cbc(key_file, iv_file, plain_file, cipher_file):
    key = read_binary_file_(key_file) 
    iv = read_binary_file_(iv_file) 
    plaintext = read_binary_file_(plain_file) 
    ciphertext_hex = read_binary_file_(cipher_file) 
    
    sm4_cbc = Sm4Cbc(key, iv, DO_ENCRYPT)
    ciphertext = sm4_cbc.update(plaintext)
    ciphertext += sm4_cbc.finish()
    cipher_std = ciphertext[0:len(plaintext)]
    #print("plaintext", plaintext.hex())
    #print("ciphertext", cipher_std.hex())
    assert ciphertext_hex == cipher_std, "加密失败"

def test_sm4_ctr (key_file, iv_file, plain_file, cipher_file):
    key = read_binary_file_(key_file) 
    iv = read_binary_file_(iv_file) 
    plaintext = read_binary_file_(plain_file) 
    ciphertext_hex = read_binary_file_(cipher_file) 

    sm4_enc = Sm4Ctr(key, iv)
    ciphertext = sm4_enc.update(plaintext)
    ciphertext += sm4_enc.finish()
    
    #print("plaintext", plaintext.hex())
    #print("ciphertext_hex", ciphertext_hex.hex())
    #print("ciphertext", ciphertext.hex())
    assert ciphertext_hex == ciphertext, "加密失败"

def read_binary_file_(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data

def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    parser.add_argument('mode', choices=['enc', 'dec'], default='enc', help='运行模式')
    args = parser.parse_args()
    directory = args.directory
    mode = args.mode
    print (f"模式:{mode}\n")

    #iv_mode = ['iv_cbc_*', 'iv_cfb_*', 'iv_ctr_*', 'iv_ecb_*', 'iv_ofb_*']
    #key_mode = ['key_cbc_*', 'key_cfb_*', 'key_ctr_*', 'key_ecb_*', 'key_ofb_*']
    #cipher_mode = ['cipher_cbc_*', 'cipher_cfb_*', 'cipher_ctr_*', 'cipher_ecb_*', 'cipher_ofb_*']
    #plain_mode = ['plain_cbc_*', 'plain_cfb_*', 'plain_ctr_*', 'plain_ecb_*', 'plain_ofb_*']

    iv_mode = ['iv_cbc_', 'iv_cfb_', 'iv_ctr_', 'iv_ecb_', 'iv_ofb_']
    key_mode = ['key_cbc_', 'key_cfb_', 'key_ctr_', 'key_ecb_', 'key_ofb_']
    cipher_mode = ['cipher_cbc_', 'cipher_cfb_', 'cipher_ctr_', 'cipher_ecb_', 'cipher_ofb_']
    plain_mode = ['plain_cbc_', 'plain_cfb_', 'plain_ctr_', 'plain_ecb_', 'plain_ofb_']
    replain_mode = ['replain_cbc_', 'replain_cfb_', 'replain_ctr_', 'replain_ecb_', 'replain_ofb_']
    enc_output_flie = ['sm4_cbc加密标准数据.txt', 'sm4_cfb加密标准数据.txt','sm4_ctr加密标准数据.txt','sm4_ecb加密标准数据.txt','sm4_ofb加密标准数据.txt']
    dec_output_flie = ['sm4_cbc解密标准数据.txt', 'sm4_cfb解密标准数据.txt','sm4_ctr解密标准数据.txt','sm4_ecb解密标准数据.txt','sm4_ofb解密标准数据.txt']

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    for j in range(5):
        # 创建输出内容
        merged_content = ""
        
        # 遍历0-9，合并对应的公钥和私钥文件
        for i in range(10):
            iv_file = os.path.join(directory, f'{iv_mode[j]}{i}')
            key_file = os.path.join(directory, f'{key_mode[j]}{i}')
            cipher_file = os.path.join(directory, f'{cipher_mode[j]}{i}')
            if mode == 'enc' :  #加密
                plain_file = os.path.join(directory, f'{plain_mode[j]}{i}')
            else:
                plain_file = os.path.join(directory, f'{replain_mode[j]}{i}')

            match j:
                case 0:
                    test_sm4_cbc(key_file, iv_file, plain_file, cipher_file)
                case 2:
                    test_sm4_ctr(key_file, iv_file, plain_file, cipher_file)
            
            # 读取文件内容
            if j != 3 : 
                iv_content = read_binary_file(iv_file)
            key_content = read_binary_file(key_file)
            cipher_content = read_binary_file(cipher_file)
            plain_content = read_binary_file(plain_file)
            length = int(len(plain_content) / 2)
            
            # 添加到合并内容
            merged_content += f"密钥=\n{key_content}\n"
            if j != 3 : 
                merged_content += f"IV=\n{iv_content}\n"
            merged_content += f"明文长度= {length:08x}\n"
            if mode == 'enc':
                merged_content += f"明文=\n{plain_content}\n"
                merged_content += f"密文=\n{cipher_content}\n"
            else:
                merged_content += f"密文=\n{cipher_content}\n"
                merged_content += f"明文=\n{plain_content}\n"
            merged_content += "\n"
        
    
        # 写入到txt文件
        if 'enc' == mode:
            output_file = enc_output_flie[j]
        else:
            output_file = dec_output_flie[j]

        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(merged_content)
        
        print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
